import numpy as np
import librosa
from audiomentations import Compose, AddGaussianNoise, PitchShift, TimeStretch
from typing import Optional

## 数据增强
# 高斯噪声添加
# 音高偏移与时间拉伸
# 随机裁剪与时间遮蔽
# 支持概率控制各增强方法

class AudioAugmenter:
    def __init__(
            self,
            sample_rate: int = 44100,
            noise_prob: float = 0.5,
            pitch_prob: float = 0.5,
            stretch_prob: float = 0.3
    ):
        """音频数据增强处理器"""
        self.sample_rate = sample_rate
        self.augment = Compose([
            AddGaussianNoise(
                min_amplitude=0.001,
                max_amplitude=0.015,
                p=noise_prob
            ),
            PitchShift(
                min_semitones=-2,
                max_semitones=2,
                p=pitch_prob
            ),
            TimeStretch(
                min_rate=0.8,
                max_rate=1.2,
                p=stretch_prob
            )
        ])

    def apply(
            self,
            y: np.ndarray,
            sr: Optional[int] = None
    ) -> np.ndarray:
        """应用增强处理"""
        sr = sr or self.sample_rate
        return self.augment(samples=y, sample_rate=sr)

    def random_crop(
            self,
            y: np.ndarray,
            target_length: int = 3,
            sr: Optional[int] = None
    ) -> np.ndarray:
        """随机裁剪音频片段"""
        sr = sr or self.sample_rate
        if len(y) <= target_length * sr:
            return y

        start = np.random.randint(0, len(y) - target_length * sr)
        return y[start:start + target_length * sr]

    def time_mask(
            self,
            y: np.ndarray,
            max_mask_length: float = 0.1  # seconds
    ) -> np.ndarray:
        """时间遮蔽增强"""
        mask_length = int(np.random.uniform(0, max_mask_length) * self.sample_rate
        if mask_length == 0:
            return y

        start = np.random.randint(0, len(y) - mask_length)
        y[start:start + mask_length] = 0
        return y

## 使用示例
# from utils.augmentation import AudioAugmenter
#
# augmenter = AudioAugmenter()
# augmented_audio = augmenter.apply(y, sr)